﻿using System.Collections.Concurrent;

namespace XServer.Utility.Pools
{
    public class AdvancedThreadPool : IDisposable
    {
        // 基础功能实现
        private readonly Queue<ITask> _taskQueue = new();
        private int _initialThreadCount;
        private volatile bool _isRunning = true;
        private int ActiveThreadCount = 0;
        private AutoResetEvent threadEvent = new AutoResetEvent(true);
        protected object locker = new object();

        // 批处理支持
        private readonly ConcurrentDictionary<Guid, BatchContext> _activeBatches = new();

        public void Start(int initialThreadCount)
        {
            _initialThreadCount = initialThreadCount;
            StartDispatcher();
        }

        private void StartDispatcher()
        {
            if (_initialThreadCount < 1) _initialThreadCount = 1;
            if (_initialThreadCount > 1000) _initialThreadCount = 1000;

            _isRunning = true;

            for (int i = 0; i < _initialThreadCount; i++)
            {
                ThreadPool.QueueUserWorkItem(new WaitCallback(ExecuteTask));
            }
            while (ActiveThreadCount < _initialThreadCount)
            {
                Thread.Sleep(100);
            }
        }

        private void ExecuteTask(object? state)
        {
            try
            {
                ActiveThreadCount = Interlocked.Increment(ref ActiveThreadCount);
                while (_isRunning)
                {
                    if (_taskQueue.Count == 0)
                    {
                        threadEvent.WaitOne();
                        continue;
                    }
                    ITask task;
                    lock (locker)
                    {
                        task = _taskQueue.Dequeue();
                    }
                    //处理
                    task.Execute();
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine($"Task execution failed: {ex}");
            }
            finally
            {
                ActiveThreadCount = Interlocked.Decrement(ref ActiveThreadCount);
            }
        }

        /// <summary>
        /// 基础任务提交
        /// </summary>
        /// <param name="task"></param>
        public void EnqueueTask(ITask task)
        {
            if (!_isRunning)
            {
                return;
            }

            lock (locker)
            {
                _taskQueue.Enqueue(task);
            }
            threadEvent.Set();
        }

        /// <summary>
        /// 批处理任务提交
        /// </summary>
        /// <param name="tasks"></param>
        /// <param name="callback"></param>
        /// <returns></returns>
        public Guid EnqueueBatch(IEnumerable<ITask> tasks, Action<BatchResult>? callback = null)
        {
            var batchId = Guid.NewGuid();
            var context = new BatchContext(tasks, callback, this);
            _activeBatches.TryAdd(batchId, context);

            // 提交批处理子任务
            foreach (var task in tasks)
            {
                var wrappedTask = new BatchSubTask(task, batchId, this);
                EnqueueTask(wrappedTask);
            }
            return batchId;
        }

        public void BatchAfterNotify(Guid id, bool success, Exception? error)
        {
            if (_activeBatches.TryGetValue(id, out var context))
            {
                context.RecordCompletion(id, success, error);
            }
        }

        public void Dispose()
        {
            _isRunning = false;
        }

        // 批处理上下文
        private class BatchContext
        {
            private readonly CountdownEvent _countdown;
            private readonly List<Exception> _errors = new();
            private readonly Action<BatchResult>? _callback;

            public int TotalTasks { get; }
            public int CompletedTasks => TotalTasks - _countdown.CurrentCount;
            private AdvancedThreadPool _outClass;

            public BatchContext(IEnumerable<ITask> tasks, Action<BatchResult>? callback, AdvancedThreadPool outClass)
            {
                _outClass = outClass;
                TotalTasks = (tasks as ICollection<ITask>)?.Count ?? 0;
                _countdown = new CountdownEvent(TotalTasks);
                _callback = callback;
            }

            public void RecordCompletion(Guid id, bool success, Exception? error)
            {
                if (!success && error != null) _errors.Add(error);
                _countdown.Signal();

                if (_countdown.IsSet)
                {
                    _callback?.Invoke(new BatchResult(
                        TotalTasks,
                        TotalTasks - _errors.Count,
                        _errors.Count > 0 ? new AggregateException(_errors) : null
                    ));

                    if (_outClass._activeBatches.TryGetValue(id, out var value))
                    {
                        _outClass._activeBatches.TryRemove(new KeyValuePair<Guid, BatchContext>(id, value));
                    }
                }
            }
        }

        // 批处理子任务包装器
        private class BatchSubTask : ITask
        {
            private readonly ITask _originalTask;
            private readonly Guid _batchId;
            private AdvancedThreadPool _outClass;

            public BatchSubTask(ITask task, Guid batchId, AdvancedThreadPool outClass)
            {
                _originalTask = task;
                _batchId = batchId;
                _outClass = outClass;
            }

            public bool Execute()
            {
                try
                {
                    _originalTask.Execute();
                    NotifyBatch(true, null);
                    return true;
                }
                catch (Exception ex)
                {
                    NotifyBatch(false, ex);
                    return false;
                }
            }

            private void NotifyBatch(bool success, Exception? error)
            {
                _outClass.BatchAfterNotify(_batchId, success, error);
            }
        }
    }

    // 批处理结果对象
    public record BatchResult(
        int TotalTasks,
        int SucceededTasks,
        AggregateException? Error
    );
}
